import random
from typing import Dict, List, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Normal
import socket
import json
import os
import pickle
import sys

class ReplayBuffer:
    """A simple numpy replay buffer."""

    def __init__(self, obs_dim: int, size: int, batch_size: int = 32):
        """Initialize."""
        self.obs_buf = np.zeros([size, obs_dim], dtype=np.float32)
        self.next_obs_buf = np.zeros([size, obs_dim], dtype=np.float32)
        self.acts_buf = np.zeros([size], dtype=np.float32)
        self.rews_buf = np.zeros([size], dtype=np.float32)
        self.done_buf = np.zeros([size], dtype=np.float32)
        self.max_size, self.batch_size = size, batch_size
        self.ptr, self.size, = 0, 0

    def store(self,
        obs: np.ndarray,
        act: np.ndarray, 
        rew: float, 
        next_obs: np.ndarray, 
        done: bool,
    ):
        """Store the transition in buffer."""
        self.obs_buf[self.ptr] = obs
        self.next_obs_buf[self.ptr] = next_obs
        self.acts_buf[self.ptr] = act
        self.rews_buf[self.ptr] = rew
        self.done_buf[self.ptr] = done
        self.ptr = (self.ptr + 1) % self.max_size
        self.size = min(self.size + 1, self.max_size)

    def sample_batch(self) -> Dict[str, np.ndarray]:
        """Randomly sample a batch of experiences from memory."""
        idxs = np.random.choice(self.size, size=self.batch_size, replace=False)
        return dict(obs=self.obs_buf[idxs],
                    next_obs=self.next_obs_buf[idxs],
                    acts=self.acts_buf[idxs],
                    rews=self.rews_buf[idxs],
                    done=self.done_buf[idxs])

    def __len__(self) -> int:
        return self.size
    

def init_layer_uniform(layer: nn.Linear, init_w: float = 3e-3) -> nn.Linear:
    """Init uniform parameters on the single layer."""
    layer.weight.data.uniform_(-init_w, init_w)
    layer.bias.data.uniform_(-init_w, init_w)

    return layer


class Actor(nn.Module):
    def __init__(
        self, 
        state_dim: int, 
        action_dim: int,
        log_std_min: float = -5,
        log_std_max: float = 2,
        hidden_dim: int = 64,
        max_action: float = 1.0,
    ):
        """Initialize."""
        super(Actor, self).__init__()
        self.trunk = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
        )
        # with separate layers works better than with Linear(hidden_dim, 2 * action_dim)
        self.mu = nn.Linear(hidden_dim, action_dim)
        self.log_sigma = nn.Linear(hidden_dim, action_dim)

        # init as in the EDAC paper
        for layer in self.trunk[::2]:
            torch.nn.init.constant_(layer.bias, 0.1)

        torch.nn.init.uniform_(self.mu.weight, -1e-3, 1e-3)
        torch.nn.init.uniform_(self.mu.bias, -1e-3, 1e-3)
        torch.nn.init.uniform_(self.log_sigma.weight, -1e-3, 1e-3)
        torch.nn.init.uniform_(self.log_sigma.bias, -1e-3, 1e-3)

        self.action_dim = action_dim
        self.max_action = max_action

        # set the log std range
        self.log_std_min = log_std_min
        self.log_std_max = log_std_max
        
    def forward(self, state: torch.Tensor) -> torch.Tensor:
        """Forward method implementation."""
        hidden = self.trunk(state)
        mu, log_sigma = self.mu(hidden), self.log_sigma(hidden)

        # clipping params from EDAC paper, not as in SAC paper (-20, 2)
        log_sigma = torch.clip(log_sigma, self.log_std_min, self.log_std_max)
        policy_dist = Normal(mu, torch.exp(log_sigma))

        # if train, deterministic is true, if test is false
        if False:
            action = mu # for test
        else:
            action = policy_dist.rsample() # for train

        tanh_action, log_prob = torch.tanh(action), None
        if True:
            # change of variables formula (SAC paper, appendix C, eq 21)
            log_prob = policy_dist.log_prob(action).sum(axis=-1)
            log_prob = log_prob - torch.log(1 - tanh_action.pow(2) + 1e-6).sum(axis=-1)
            log_prob = log_prob.sum(-1, keepdim=True) # 参考原来的SAC2 这里进行sum操作
        return tanh_action * self.max_action, log_prob
        
    
class CriticQ(nn.Module):
    def __init__(self, in_dim: int):
        
        """Initialize."""
        super(CriticQ, self).__init__()
        self.trunk = nn.Sequential(
            nn.Linear(in_dim, hidden_num),
            nn.ReLU(),
            nn.Linear(hidden_num, hidden_num),
            nn.ReLU(),
            nn.Linear(hidden_num, hidden_num),
            nn.ReLU(),
        )
        for layer in self.trunk[::2]:
            torch.nn.init.constant_(layer.bias, 0.1)
            
        # self.hidden1 = nn.Linear(in_dim, hidden_num)
        self.out = nn.Linear(hidden_num, 1)
        self.out = init_layer_uniform(self.out)

    def forward(
        self, state: torch.Tensor, action: torch.Tensor
    ) -> torch.Tensor:
        """Forward method implementation."""
        x = torch.cat((state, action), dim=-1)
        # x = F.relu(self.hidden1(x))
        x = self.trunk(x)
        value = self.out(x)
        
        return value
    
    
class CriticV(nn.Module):
    def __init__(self, in_dim: int):
        """Initialize."""
        super(CriticV, self).__init__()
        self.trunk = nn.Sequential(
            nn.Linear(in_dim, hidden_num),
            nn.ReLU(),
            nn.Linear(hidden_num, hidden_num),
            nn.ReLU(),
            nn.Linear(hidden_num, hidden_num),
            nn.ReLU(),
        )
        for layer in self.trunk[::2]:
            torch.nn.init.constant_(layer.bias, 0.1)
            
        self.hidden1 = nn.Linear(in_dim, hidden_num)
        self.out = nn.Linear(hidden_num, 1)
        self.out = init_layer_uniform(self.out)

    def forward(self, state: torch.Tensor) -> torch.Tensor:
        """Forward method implementation."""
        # x = F.relu(self.hidden1(state))
        x = self.trunk(state)
        value = self.out(x)
        
        return value
    
class SACAgent:
    """SAC agent interacting with environment.
    
    Attrtibutes:
        actor (nn.Module): actor model to select actions
        actor_optimizer (Optimizer): optimizer for training actor
        vf (nn.Module): critic model to predict state values
        vf_target (nn.Module): target critic model to predict state values
        vf_optimizer (Optimizer): optimizer for training vf
        qf_1 (nn.Module): critic model to predict state-action values
        qf_2 (nn.Module): critic model to predict state-action values
        qf_1_optimizer (Optimizer): optimizer for training qf_1
        qf_2_optimizer (Optimizer): optimizer for training qf_2
        memory (ReplayBuffer): replay memory
        batch_size (int): batch size for sampling
        gamma (float): discount factor
        tau (float): parameter for soft target update
        initial_random_steps (int): initial random action steps
        policy_update_freq (int): policy update frequency
        device (torch.device): cpu / gpu
        target_entropy (int): desired entropy used for the inequality constraint
        log_alpha (torch.Tensor): weight for entropy
        alpha_optimizer (Optimizer): optimizer for alpha
        transition (list): temporory storage for the recent transition
        total_step (int): total step numbers
        is_test (bool): flag to show the current mode (train / test)
    """
    
    def __init__(
        self,
        memory_size: int,
        batch_size: int,
        gamma: float = 0.99,
        tau: float = 5e-3,
        initial_random_steps: int = int(1e4),
        policy_update_freq: int = 2,
        obs_dim: int = 4,
        action_dim: int = 1,
        q_num: int =50,
    ):
        """Initialize."""
        obs_dim = obs_dim
        action_dim = action_dim

        self.memory = ReplayBuffer(obs_dim, memory_size, batch_size)
        self.batch_size = batch_size
        self.gamma = gamma
        self.tau = tau
        self.initial_random_steps = initial_random_steps
        self.policy_update_freq = policy_update_freq

        # device: cpu / gpu
        self.device = torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )
        print(self.device)
        
        # automatic entropy tuning
        self.target_entropy = -np.prod((action_dim,)).item()  # heuristic
        self.log_alpha = torch.zeros(1, requires_grad=True, device=self.device)
        self.alpha_optimizer = optim.Adam([self.log_alpha], lr=3e-4)

        # actor
        self.actor = Actor(obs_dim, action_dim).to(self.device)
        
        # v function
        self.vf = CriticV(obs_dim).to(self.device)
        self.vf_target = CriticV(obs_dim).to(self.device)
        self.vf_target.load_state_dict(self.vf.state_dict())
        
        # q function, N for overestimation
        self.q_model = list()
        self.q_model_optim = list()
        for i in range(q_num):
            self.q_model.append(CriticQ(obs_dim + action_dim).to(self.device))
        for i in range(q_num):
            self.q_model_optim.append(optim.Adam(self.q_model[i].parameters(), lr=3e-4))

        # optimizers
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=3e-4)
        self.vf_optimizer = optim.Adam(self.vf.parameters(), lr=3e-4)

        # transition to store in memory
        self.transition = list()
        
        # total steps count
        self.total_step = 0

        # network socket
        REMOTE_HOST = '127.0.0.1'
        REMOTE_PORT = 10905
        self.tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.tcp_socket.connect((REMOTE_HOST, REMOTE_PORT))
        
        self.rudder = 0
    
    def select_action(self, state: np.ndarray) -> np.ndarray:
        """Select an action from the input state."""
        # if initial random action should be conducted
        if self.total_step < self.initial_random_steps:
            selected_action = np.array(random.uniform(-1, 1))
        else:
            selected_action = self.actor(
                torch.FloatTensor(state).to(self.device)
            )[0].detach().cpu().numpy()
            
        self.transition = [state, selected_action]
        
        return selected_action
    
    def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool]:
        """Take an action and return the response of the env."""
        if isinstance(action, float):
            action = action
        else:
            action = action.item()
        action_data, rudder = encode_data(action, rudder = self.rudder, reset_flag=0)
        self.rudder = rudder
        self.tcp_socket.send(action_data)# action
        info, addr = self.tcp_socket.recvfrom(1024)
        observation, reward, terminated, truncated = decode_data(info)# next state
        reward = torch.tensor(reward, device=self.device).unsqueeze(0)# reward
        done = terminated or truncated
        done = torch.tensor(done, device=self.device).unsqueeze(0)
        next_state = torch.tensor(observation, dtype=torch.float32, device=self.device).unsqueeze(0)

        self.transition += [reward, next_state, done]
        self.memory.store(*self.transition)
        
        # memory save for pretrain
        memory_demo.append(tuple(self.transition))
    
        return next_state, reward, done
    
    def update_model(self) -> Tuple[torch.Tensor, ...]:
        """Update the model by gradient descent."""
        device = self.device  # for shortening the following lines
        
        samples = self.memory.sample_batch()
        state = torch.FloatTensor(samples["obs"]).to(device)
        next_state = torch.FloatTensor(samples["next_obs"]).to(device)
        action = torch.FloatTensor(samples["acts"].reshape(-1, 1)).to(device)
        reward = torch.FloatTensor(samples["rews"].reshape(-1, 1)).to(device)
        done = torch.FloatTensor(samples["done"].reshape(-1, 1)).to(device)
        new_action, log_prob = self.actor(state)
        
        # alpha loss (dual problem)
        alpha_loss = (
            -self.log_alpha.exp() * (log_prob + self.target_entropy).detach()
        ).mean()

        self.alpha_optimizer.zero_grad()
        alpha_loss.backward()
        self.alpha_optimizer.step()
        
        alpha = self.log_alpha.exp()  # used for the actor loss calculation
        
        # q function loss
        mask = 1 - done
        # 2Q change to N-Q-network
        q_pred_l = list()
        for i in range(len(self.q_model)):
            q_pred_l.append(self.q_model[i](state, action))

        v_target = self.vf_target(next_state)
        q_target = reward + self.gamma * v_target * mask
        
        q_loss_l = list()
        for i in range(len(self.q_model)):
            q_loss_l.append( F.mse_loss(q_pred_l[i], q_target.detach()))

        qf_loss = torch.tensor([0.0])
        # train Q functions
        for i in range(len(self.q_model)):
            self.q_model_optim[i].zero_grad()
            q_loss_l[i].backward()
            self.q_model_optim[i].step()
            qf_loss += q_loss_l[i]
        
        # v function loss
        v_pred = self.vf(state) 
        for i in range(len(self.q_model)):
            if i == 0:
                q_pred = self.q_model[i](state, new_action)
            else:
                q_pred = torch.min(q_pred, self.q_model[i](state, new_action))

        v_target = q_pred - alpha * log_prob
        vf_loss = F.mse_loss(v_pred, v_target.detach())
        
        if self.total_step % self.policy_update_freq == 0:
            # actor loss
            advantage = q_pred - v_pred.detach()
            actor_loss = (alpha * log_prob - advantage).mean()
        
            self.actor_optimizer.zero_grad()
            actor_loss.backward()
            self.actor_optimizer.step()
        
            # target update (vf)
            self._target_soft_update()
        else:
            actor_loss = torch.zeros(1)
            
        # train V function
        self.vf_optimizer.zero_grad()
        vf_loss.backward()
        self.vf_optimizer.step()
        
        return actor_loss.data, qf_loss.data, vf_loss.data, alpha_loss.data
    
    def train(self, num_frames: int):
        """Train the agent."""
        # env reset
        action_data, rudder = encode_data(0.0, rudder=self.rudder, reset_flag=1)
        self.rudder = rudder
        # print(data)
        self.tcp_socket.send(action_data)
        # Initialize the environment and get its state
        info, addr = self.tcp_socket.recvfrom(1024)
        observation, _, _, _= decode_data(info)
        state = torch.tensor(observation, dtype=torch.float32, device=self.device).unsqueeze(0)

        score = 0
        i_episode = 0
        # frame_count = 0
        for self.total_step in range(1, num_frames + 1):
            action = self.select_action(state)
            next_state, reward, done = self.step(action)

            state = next_state
            
            # frame_count += 1
            # score = (score + reward)/ frame_count
            score += reward

            # if episode ends
            if done:
                i_episode += 1
                recorder['score'].append(score.squeeze(0).cpu().numpy())
                print(
                    f'episode is {i_episode}, score is {score.squeeze(0).cpu().numpy()}')

                # state = env.reset()
                action = random.uniform(-1.0, 1.0)
                action_data, rudder = encode_data(action, rudder=0, reset_flag=1)
                self.rudder = rudder
                # print(data)
                self.tcp_socket.send(action_data)
                # Initialize the environment and get its state
                info, addr = self.tcp_socket.recvfrom(1024)
                observation, reward, terminated, truncated= decode_data(info)
                state = torch.tensor(observation, dtype=torch.float32, device=self.device).unsqueeze(0)
                score = 0
                # frame_count = 0

            # if training is ready
            if (
                len(self.memory) >= self.batch_size 
                and self.total_step > self.initial_random_steps
            ):
                losses = self.update_model()
                recorder['actor_losses'].append(losses[0])
                recorder['qf_losses'].append(losses[1])
                recorder['vf_losses'].append(losses[2])
                recorder['alpha_losses'].append(losses[3])

       
    def _target_soft_update(self):
        """Soft-update: target = tau*local + (1-tau)*target."""
        tau = self.tau
        
        for t_param, l_param in zip(
            self.vf_target.parameters(), self.vf.parameters()
        ):
            t_param.data.copy_(tau * l_param.data + (1.0 - tau) * t_param.data)
    

def save_best_model(path):

    #gpu tensor format save
    torch.save(obj={
        'actor_state_dict': agent.actor.state_dict(),
        'num_frames' : num_frames,
        'memory_size' : memory_size,
        'batch_size' : batch_size,
        'initial_random_steps' : initial_random_steps,
        }, f=path)


def decode_data(data):
    # load json
    recv_obser = json.loads(data)
    # print(recv_obser)
    observation = []
    observation_dict = recv_obser['observation']
    for key in observation_dict:
        observation.append(observation_dict[key])
    reward = recv_obser['reward']
    terminated = recv_obser['terminated']
    truncated = recv_obser['truncated']
    return observation, reward, terminated, truncated

def encode_data(action, rudder, reset_flag=1):

    # rudder = rudder + float(action) * 5 # add?
    rudder = action * 30

    if rudder >= 30.0:
        rudder = 30.0
    elif rudder <= -30.0:
        rudder = -30.0
    
    origin_data = {'boatname':'SLM7001',
                   'restart': reset_flag,
                   'rudl': rudder,
                   'rudr': rudder,
                   'rspl': 1000,
                   'rspr': 1000,
                   'subSystem': "control"
                   }

    data = json.dumps(origin_data, sort_keys=True, indent=4, separators=(',', ':'))
    return data.encode('utf-8'), rudder


HOME = os.getcwd()
# 训练效果记录
recorder = dict({
    'score': [],
    'actor_losses': [],
    'qf_losses': [],
    'vf_losses': [],
    'alpha_losses': []
})

# parameters
num_frames = 400_000 # review plot, whether add frames and memory?
memory_size = 100_000
batch_size = 128
initial_random_steps = 10_000
hidden_num = 128
q_ensemble = 10

memory_demo = list()
save_model_path = os.path.join(HOME, f'SAC/sac_{q_ensemble}_model_add.pt')
recorder_path = os.path.join(HOME, f'SAC/sac_{q_ensemble}_record_add.pkl')
save_memory_path = os.path.join(HOME, f'SAC/sac_{q_ensemble}_pretrain_memory.pkl')

agent = SACAgent(
    memory_size=memory_size, batch_size=batch_size, initial_random_steps=initial_random_steps, obs_dim=3, action_dim=1, q_num=q_ensemble
)

agent.train(num_frames)

# save to file for ploting
with open(recorder_path, 'wb') as f:
    pickle.dump(recorder, f) 
    
# save memory for pretrained
with open(save_memory_path, 'wb') as f:
    pickle.dump(memory_demo, f)
    
save_best_model(save_model_path)
print('Complete')